import json
import multiprocessing
import os
import torch
from torch import nn
from d2l import torch as d2l

batch_size = 512
max_len = 128
num_workers = d2l.get_dataloader_workers()
data_dir = d2l.download_extract('SNLI')
train_set = d2l.SNLIDataset(d2l.read_snli(data_dir,True),max_len)
test_set = d2l.SNLIDataset(d2l.read_snli(data_dir,False),max_len)
train_iter = torch.utils.data.DataLoader(train_set,batch_size,
                                         shuffle=True,num_workers=num_workers)
test_iter = torch.utils.data.DataLoader(test_set,batch_size,
                                        num_workers=num_workers)

class BERTClassifier(nn.Module):
    def __init__(self,bert):
        super.__Init__()
        self.encoder = bert.encoder
        self.hidden = bert.hidden
        self.output = nn.Linear(256,3)
    
    def forward(self,inputs):
        tokens_X,segments_X,valid_lens_X = inputs
        encoded_X = self.encoder(tokens_X,segments_X,valid_lens_X)
        return self.output(self.hidden(encoded_X[:,0,:]))
